#include <iostream>
#include <time.h>
#include<thrust/host_vector.h>
#include <thrust/device_vector.h>
#include<thrust/sort.h>
#include<thrust/iterator/counting_iterator.h>
#include<thrust/binary_search.h>
#include <stdio.h>
#include "CudaRand.h"
using namespace std;

#define POP_SIZE 2000000
#define LOCATION 1000
#define INIT_INFECTION 50
#define INFECT_RATE 0.1
#define BRANCH_RATE 0.2
#define ITER 5

__global__ void Sum2(int* sum){

	__shared__ int local[1];
	
	int tid_in_block = threadIdx.x * blockDim.y + threadIdx.y;
	
	if(tid_in_block == 0){
		local[0] = 0;
	}
		
	__syncthreads();
		
	for(int i = 0; i < 1000; i++){
		atomicAdd(&(local[0]),1);
	}
		
	__syncthreads();
		
		
	if(tid_in_block == 0){
		atomicAdd(&(sum[0]),local[0]);
	}
}

__global__ void Sum1(int* sum){
	for(int i = 0; i < 1000; i++) atomicAdd(&(sum[0]),1);
}

__device__ void EnterCS(int* mutex){
        do{}while(atomicCAS(&(mutex[0]),0,1));
		//printf("(%d,%d) with size  %d inside CS \n",threadIdx.x,blockIdx.x,gridDim.x);

};

__device__ void QuitCS(int* mutex){
        __threadfence();
		mutex[0] = 0;
};

__device__ double f(){
	double x = 0;
	for(int i = 0; i < 100; i++){
		x += sin((double)i);
	}
	return x;
}

__device__ double g(){
	double y = 0;
	for(int i = 0; i < 100; i++){
		y += cos((double)i);
	}
	return y;
}

struct CastA{
	__device__
	int operator()(RandGen& rand_obj){
		double result = 0;
	
		for(int i = 0; i < 100; i++){

		switch(rand_obj.cudaRandInt()%3){
		case 0:
			result = f();
			break;
		case 1:
			result = g();
			break;
		case 2:
			result += f();
			result += g();
			break;
		}
		
		}
		return (int)result;
	}
};

struct CastB{
	__device__
	int operator()(RandGen& rand_obj){
		double result = 0;
		
		for(int i = 0; i < 100; i++){
		result = f();
		result = g();
		result += f();
		result += g();
		}
		
		return (int)result;
	}
};

__global__ void critical_test(int* a){
	int index = blockDim.x*blockIdx.x+threadIdx.x;
	EnterCS(a);
	printf("thread %d is in section 1\n",index);
	QuitCS(a);
	
	// EnterCS(a);
	// printf("thread %d is in section 2\n",index);
	// QuitCS(a);
	
}

__device__ void FurtherDeter(int* result, RandGen& rand_obj, int idx){
	if(rand_obj.cudaRandUni() < BRANCH_RATE){
		result[idx] = 1;
	}
}

__global__ void InfectSimple(int* location_begin, int* location_end, int* status, int* result, RandGen* rand_vec){
	int idx = blockIdx.x * blockDim.x + threadIdx.x;
	int start_idx = location_begin[idx];
	int end_idx = location_end[idx];
	RandGen& rand_obj = rand_vec[idx];
	
	for(int i = start_idx; i < end_idx; i++){
		if(status[i] == 1){
			for(int j = start_idx; j < end_idx; j++){
				if(status[j] == 0 && rand_obj.cudaRandUni() < INFECT_RATE) FurtherDeter(result,rand_obj,j);
			}
		}
	}
}

// each block assign to a location
__global__ void Infect1(int* location_begin, int* location_end, int* status, int* result, RandGen* rand_vec){
	int tid = threadIdx.x;
	int bid = blockIdx.x;
	int idx = tid + bid * blockDim.x;
	
	int block_start_idx = location_begin[bid];
	int block_end_idx = location_end[bid];
	
	int location_len = block_end_idx - block_start_idx;
	int thread_len = location_len / blockDim.x + (location_len % blockDim.x != 0);
	int local_start_idx = block_start_idx  + tid * thread_len;
	int local_end_idx = local_start_idx + thread_len;
	local_end_idx = local_end_idx < block_end_idx ? local_end_idx : block_end_idx;
	
	RandGen& rand_obj = rand_vec[idx];
	
	for(int i = local_start_idx; i < local_end_idx; i++){
		if(status[i] == 1){
			for(int j = block_start_idx; j < block_end_idx; j++){
				if(status[j] == 0 && rand_obj.cudaRandUni() < INFECT_RATE) FurtherDeter(result,rand_obj,j);
			}
		}
	}
	
}
	
__global__ void SetPos1(int* location_begin, int* location_end, int* pos_vec){
	int idx = blockIdx.x * blockDim.x + threadIdx.x;
	int start_idx = location_begin[idx];
	int end_idx = location_end[idx];
	
	for(int i = start_idx; i < end_idx; i++){
		pos_vec[i] = idx;
	}
}

__global__ void SetPos2(int* pos_vec, int* thread_pos, int thread_len){
	int idx = blockIdx.x * blockDim.x + threadIdx.x;
	int pos = thread_len * idx;
	
	if(pos < POP_SIZE) thread_pos[idx] = pos_vec[pos];
	
}

__global__ void Infect2(int* location_begin, int* location_end, int* status, int* result, RandGen* rand_vec,
	int* thread_pos, int thread_len){
	
	int idx = blockIdx.x * blockDim.x + threadIdx.x;
	int location_idx = thread_pos[idx];
	RandGen& rand_obj = rand_vec[idx];
	
	if(location_idx != POP_SIZE){
	
	int block_start = location_begin[location_idx];
	int block_end = location_end[location_idx];
	
	int thread_start = idx * thread_len;
	int thread_end = (idx + 1) * thread_len;
	thread_end = thread_end < block_end ? thread_end : block_end;
	
	for(int i = thread_start; i < thread_end; i++){
		if(status[i] == 1){
			for(int j = block_start; j < block_end; j++){
				if(status[j] == 0 && rand_obj.cudaRandUni() < INFECT_RATE) FurtherDeter(result,rand_obj,j);
			}
		}
	}
	
	}
	
}

int main(){	
	cudaError err;
	time_t start;
	
	srand(time(NULL));
	
	thrust::host_vector<int> h_agent_key; // index of each agent
	thrust::host_vector<int> h_location; // the location of the agent go
	thrust::host_vector<int> h_status(POP_SIZE,0);
	thrust::host_vector<RandGen> h_rand_vec;
	
	for(int i = 0; i < POP_SIZE; i++){
		h_agent_key.push_back(i);
		h_location.push_back(((int)(cos((double)i)*LOCATION))%LOCATION);
		//h_location.push_back(rand()%LOCATION);
		h_rand_vec.push_back(RandGen(rand(),rand(),rand(),rand()));
	}
	
	// initial infection
	for(int i = 0; i < INIT_INFECTION; i++) h_status[rand()%POP_SIZE] = 1;
	
	thrust::device_vector<int> agent_key = h_agent_key;
	thrust::device_vector<int> location = h_location;
	thrust::device_vector<RandGen> rand_vec = h_rand_vec;
	// sort agent key by location number
	thrust::sort_by_key(location.begin(),location.end(),agent_key.begin());
	
	thrust::device_vector<int> location_begin(LOCATION);
	thrust::device_vector<int> location_end(LOCATION);
	thrust::counting_iterator<int> count_location(0);
	
	thrust::lower_bound(
		location.begin(),location.end(),
		count_location,count_location+LOCATION,
		location_begin.begin());
		
	thrust::upper_bound(
		location.begin(),location.end(),
		count_location,count_location+LOCATION,
		location_end.begin());
	
	thrust::device_vector<int> status1 = h_status;
	thrust::device_vector<int> result(POP_SIZE,0);
	
	// for(int i = 0; i < LOCATION; i++){
		// cout << "location " << i << " has ppl: " << location_end[i] << " to " <<  location_begin[i] << endl;
	// }
	
	cout << "started" << endl;
	
	// start = clock();
	
	// for(int i = 0; i < ITER; i++){
	
	// cout << "iter " << i << endl;
	
	// thrust::fill(result.begin(),result.end(),0);
	
	// InfectSimple<<<LOCATION/100,100>>>(
		// thrust::raw_pointer_cast(location_begin.data()),
		// thrust::raw_pointer_cast(location_end.data()),
		// thrust::raw_pointer_cast(status1.data()),
		// thrust::raw_pointer_cast(result.data()),
		// thrust::raw_pointer_cast(rand_vec.data()));
	// err = cudaThreadSynchronize();
	// if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	// thrust::copy(result.begin(),result.end(),status1.begin());
	
	// }
	
	// cout << "intial infection: " << INIT_INFECTION << " and result infected: " << thrust::reduce(result.begin(),result.end()) << endl;
	// cout << "time take for part 1 is " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	
/**************** part 2 *******************************/
	thrust::device_vector<int> status2 = h_status;
	
	start = clock();
	
	for(int i = 0 ; i < ITER; i++){
	
	thrust::fill(result.begin(),result.end(),0);
	
	Infect1<<<LOCATION,100>>>(
		thrust::raw_pointer_cast(location_begin.data()),
		thrust::raw_pointer_cast(location_end.data()),
		thrust::raw_pointer_cast(status2.data()),
		thrust::raw_pointer_cast(result.data()),
		thrust::raw_pointer_cast(rand_vec.data()));
	err = cudaThreadSynchronize();
	if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	cout << "new infected at " << i << " is " << thrust::reduce(result.begin(),result.end()) << endl;
	
	thrust::copy(result.begin(),result.end(),status2.begin());
	
	}
	
	cout << "intial infection: " << INIT_INFECTION << " and result infected: " << thrust::reduce(result.begin(),result.end()) << endl;
	cout << "time take for part 2 is " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	
	
/****************** part 3 ************************/
	int number_threads = LOCATION * 100;
	int thread_len = POP_SIZE / number_threads + (POP_SIZE % number_threads != 0);
	thrust::device_vector<int> status3 = h_status;
	thrust::device_vector<int> pos_vec(POP_SIZE);
	thrust::device_vector<int> thread_location(number_threads,POP_SIZE);
	cout << "thread len :" << thread_len << endl;
	
	start = clock();
	for(int i = 0 ; i < ITER; i++){
	
	thrust::fill(result.begin(),result.end(),0);
	
	SetPos1<<<LOCATION/100,100>>>(
		thrust::raw_pointer_cast(location_begin.data()),
		thrust::raw_pointer_cast(location_end.data()),
		thrust::raw_pointer_cast(pos_vec.data()));
		
	SetPos2<<<LOCATION,100>>>(
		thrust::raw_pointer_cast(pos_vec.data()),
		thrust::raw_pointer_cast(thread_location.data()),
		thread_len);
	
	Infect2<<<LOCATION,100>>>(
		thrust::raw_pointer_cast(location_begin.data()),
		thrust::raw_pointer_cast(location_end.data()),
		thrust::raw_pointer_cast(status3.data()),
		thrust::raw_pointer_cast(result.data()),
		thrust::raw_pointer_cast(rand_vec.data()),
		thrust::raw_pointer_cast(thread_location.data()),
		thread_len);
	
	err = cudaThreadSynchronize();
	if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	cout << "new infected at " << i << " is " << thrust::reduce(result.begin(),result.end()) << endl;
	
	thrust::copy(result.begin(),result.end(),status3.begin());
	
	}
	cout << "intial infection: " << INIT_INFECTION << " and result infected: " << thrust::reduce(result.begin(),result.end()) << endl;
	cout << "time take for part 3 is " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	
	// thrust::device_vector<int> a(1,0);
	// Sum1<<<1000,1000>>>(thrust::raw_pointer_cast(a.data()));
	// err = cudaThreadSynchronize();
    // if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	// cout << "result is " << a[0] << " which takes time: " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
		
	// start = clock();
	// thrust::device_vector<int> b(1,0);
	// Sum2<<<1000,1000>>>(thrust::raw_pointer_cast(b.data()));
	// err = cudaThreadSynchronize();
	// if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	// cout << "result is " << b[0] << " which takes time: " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	
	
	// thrust::device_vector<int> c(1,0);
	// critical_test<<<10,10>>>(thrust::raw_pointer_cast(c.data()));
	// if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	return 0;
}
